import torch
import numpy as np
from lowpass_filter import lowpass_filter


def wdm_demultiplex(WDM_sigout, Signal):
    """
    Demultiplex the WDM signal.

    Parameters:
    WDM_sigout: The multiplexed WDM signal.
    Signal: A dictionary containing the signal parameters.

    Returns:
    WDM_sigout_X: The demultiplexed signal for X.
    WDM_sigout_Y: The demultiplexed signal for Y.
    """
    Sampling_num = Signal['Sampling_num']
    WDM_channel = Signal['WDM_channel']
    df_matrix = Signal['df_matrix']
    simu_time = Signal['simu_time']
    channel_spacing_Hz = Signal['channel_spacing_Hz']

    WDM_sigout_X = torch.zeros((Sampling_num, WDM_channel), dtype=torch.complex64)
    WDM_sigout_Y = torch.zeros((Sampling_num, WDM_channel), dtype=torch.complex64)

    for k in range(WDM_channel):
        phase_array = 2 * np.pi * df_matrix[k] * simu_time
        phase_array = phase_array.numpy() if isinstance(phase_array, torch.Tensor) else phase_array
        phase_array1 = np.exp(-1j * phase_array)
        phase_array1 = phase_array1.reshape(-1, 1)

        h = lowpass_filter(Signal, channel_spacing_Hz/2)

        WDM_sig_X = WDM_sigout[:, 0].numpy() * phase_array1.flatten()
        WDM_sig_Y = WDM_sigout[:, 1].numpy() * phase_array1.flatten()

        E_fft_X = np.fft.fftshift(np.fft.fft(WDM_sig_X))
        E_fft_X *= h
        WDM_sigout_X[:, k] = torch.from_numpy(np.fft.ifft(np.fft.ifftshift(E_fft_X)))

        E_fft_Y = np.fft.fftshift(np.fft.fft(WDM_sig_Y))
        E_fft_Y *= h
        WDM_sigout_Y[:, k] = torch.from_numpy(np.fft.ifft(np.fft.ifftshift(E_fft_Y)))

    return WDM_sigout_X, WDM_sigout_Y